#ifndef TEST_HPP
#define TEST_HPP

#include <iostream>
#include <cmath>
#include <vector>
#include <algorithm>
#include <fstream>
#include <Eigen/Sparse>
#include <Eigen/Dense>
#include "Poisson.hpp"
#include "Function.hpp"
#include "json.hpp"

struct BoundaryConditions_One_Dimension
{
    enum Type
    {
        DIRICHLET,
        NEUMANN,
        Mixed
    };
    Type type;
    std::vector<int> flag;
    std::vector<double> one_point;
    std::vector<double> value;
};

struct MGParams_One_Dimension
{
    int grid_resolution;
    enum Restriction
    {
        full_weighting,
        injection
    };
    Restriction restriction;
    int recursions;
    enum Interpolation
    {
        linear,
        quadratic
    };
    Interpolation interpolation;
    enum Type
    {
        V_Cycles,
        FMG
    };
    Type type;
    int max_iterations;
    double relative_accuracy;
    std::shared_ptr<Function_One_Dimension> initial_guess;
    int v1; // Forward Relax times
    int v2; // Backward Relax times
};

class Multigrid_One_Dimension
{
private:
    MGParams_One_Dimension MG;
    BoundaryConditions_One_Dimension BC;
    std::shared_ptr<Function_One_Dimension> source_term;
    double h;
    double h2;
    int n;
    int m;
    double omega = 0.75; // Jacobi权重系数
    Eigen::SparseMatrix<double> A;
    Eigen::VectorXd f;
    Eigen::VectorXd U;
    std::vector<int> ipiv;
    std::vector<std::vector<int>> index_map;

    Eigen::SparseMatrix<double> I_h_2h(int N)
    {
        Eigen::SparseMatrix<double> X(N / 2 - 1, N - 1);
        int flag = 0;
        if (MG.restriction == MGParams_One_Dimension::injection)
        {
            for (int i = 0; i < N / 2 - 1; i++)
            {
                X.insert(i, 2 * i) = 1;
            }
        }
        if (MG.restriction == MGParams_One_Dimension::full_weighting)
        {
            for (int i = 0; i < N / 2 - 1; i++)
            {
                X.insert(i, flag++) = 0.25;
                X.insert(i, flag++) = 0.5;
                X.insert(i, flag) = 0.25;
            }
        }
        return X;
    }

    Eigen::SparseMatrix<double> I_2h_h(int N)
    {
        Eigen::SparseMatrix<double> X(N - 1, N / 2 - 1);
        if (MG.interpolation == MGParams_One_Dimension::quadratic)
        {
            for (int i = 0; i < N / 2 - 1; i++)
            {
                int fine_idx = 2 * i + 1; // 对应的细网格起始位置

                // 二次插值模板：1/4, 1/2, 1/4
                if (fine_idx - 1 >= 0)
                {
                    X.insert(fine_idx - 1, i) = 0.25;
                }
                X.insert(fine_idx, i) = 0.5;
                if (fine_idx + 1 < N - 1)
                {
                    X.insert(fine_idx + 1, i) = 0.25;
                }
            }
        }
        if (MG.interpolation == MGParams_One_Dimension::linear)
        {
            int flag = 0;
            for (int i = 0; i < N / 2 - 1; i++)
            {
                X.insert(flag++, i) = 0.5;
                X.insert(flag++, i) = 1.0;
                X.insert(flag, i) = 0.5;
            }
        }
        return X;
    }

    void initialization()
    {
        int total_points = m + 2;
        n = m + 1;
        A.resize(total_points, total_points);
        A.reserve(Eigen::VectorXi::Constant(total_points, 3));
        U.resize(total_points);
        f.resize(total_points);
        ipiv.resize(total_points, 0);
        if (MG.type == MGParams_One_Dimension::V_Cycles)
        {
            for (int i = 0; i < total_points; ++i)
            {
                U[i] = MG.initial_guess->operator()((i)*h);
            }
        }
        std::cout << "1" << std::endl;

        if (BC.value.size() < 2)
        {
            throw std::runtime_error("BoundaryConditions_One_Dimension value vector too small");
        }
        if (BC.type == BoundaryConditions_One_Dimension::DIRICHLET)
        {
            for (int i = 0; i < total_points; ++i)
            {
                std::cout << "3" << std::endl;
                if (i == 0)
                {
                    A.insert(i, i) = 1;
                    f[i] = BC.value[0];
                }
                else if (i == total_points - 1)
                {
                    A.insert(i, i) = 1;
                    f[i] = BC.value[1];
                }
                else
                {
                    A.insert(i, i) = 2.0 / h2;
                    A.insert(i, i - 1) = A.insert(i, i + 1) = -1.0 / h2;
                    f[i] = (*source_term)((i)*h);
                }
            }
            return;
        }

        if (BC.type == BoundaryConditions_One_Dimension::NEUMANN)
        {
            for (int i = 0; i < total_points; ++i)
            {
                if (i == int(BC.one_point[0] / h))
                {
                    A.insert(i, i) = 1.0;
                    f[i] = BC.one_point[1];
                }
                else if (i == 0)
                {
                    A.insert(i, i) = 1.0 / h;
                    A.insert(i, i + 1) = -1.0 / h;
                    f[i] = -BC.value[0] + (h * (*source_term)(0)) / 2.0;
                }
                else if (i == total_points - 1)
                {
                    A.insert(i, i - 1) = -1.0 / h;
                    A.insert(i, i) = 1.0 / h;
                    f[i] = BC.value[1] + (h * (*source_term)(1)) / 2.0;
                }
                else
                {
                    A.insert(i, i) = 2.0 / h2;
                    A.insert(i, i - 1) = A.insert(i, i + 1) = -1.0 / h2;
                    f[i] = (*source_term)(i * h);
                }
            }
            return;
        }

        if (BC.type == BoundaryConditions_One_Dimension::Mixed)
        {
            if (BC.flag[0] == 0 && BC.flag[1] == 1)
            {
                A.insert(0, 0) = 1;
                f[0] = BC.value[0];
                A.insert(m + 1, m + 1) = 1.0 / h;
                A.insert(m + 1, m) = -1.0 / h;
                f[m + 1] = BC.value[1] + (h * (*source_term)(0)) / 2.0;
            }
            else if (BC.flag[0] == 1 && BC.flag[1] == 0)
            {
                A.insert(0, 0) = 1.0 / h;
                A.insert(0, 1) = -1.0 / h;
                f[0] = -BC.value[0] + (h * (*source_term)(0)) / 2.0;
                A.insert(m + 1, m + 1) = 1;
                f[m + 1] = BC.value[1];
            }
            else
            {
                std::cout << "Invalid Input for Mixed Boundary Conditions One Dimension!!!" << std::endl;
                return;
            }
            for (int i = 1; i < total_points - 1; ++i)
            {
                A.insert(i, i) = 2.0 / h2;
                A.insert(i, i - 1) = A.insert(i, i + 1) = -1.0 / h2;
                f[i] = (*source_term)(i * h);
            }
            return;
        }
    }

    // 加权 Jacobi 迭代
    void weightedJacobi(const Eigen::SparseMatrix<double> &A, Eigen::VectorXd &x,
                        const Eigen::VectorXd &b, int v)
    {
        if (A.rows() != A.cols() || A.rows() != b.size() || b.size() != x.size())
        {
            throw std::runtime_error("Matrix and vector dimensions mismatch");
        }

        // Extract diagonal and compute its inverse
        Eigen::VectorXd D_inv = A.diagonal().cwiseInverse();

        // Create matrix R = L + U (off-diagonal parts)
        Eigen::SparseMatrix<double> R = A;
        R.diagonal().setZero();

        for (int k = 0; k < v; ++k)
        {
            x = (1.0 - omega) * x + omega * D_inv.cwiseProduct(b - R * x);
        }
    }

public:
    Multigrid_One_Dimension(const MGParams_One_Dimension &mg, const BoundaryConditions_One_Dimension &bc, std::shared_ptr<Function_One_Dimension> f)
        : MG(mg), BC(bc), source_term(f)
    {
        h = 1.0 / (MG.grid_resolution + 1);
        h2 = h * h;
        m = MG.grid_resolution;
        std::cout << m << std::endl;
        initialization();
        std::cout << "2" << std::endl;
    }

    Eigen::VectorXd V_cycles(Eigen::SparseMatrix<double> A1, Eigen::VectorXd v1, Eigen::VectorXd f1)
    {

        const auto &I_h_2h_matrix = I_h_2h(f1.size() + 1);
        const auto &I_2h_h_matrix = I_2h_h(f1.size() + 1);
        std::cout << "A1:" << std::endl;
        std::cout << A1 << std::endl;
        std::cout << "I_h_2h:" << std::endl;
        std::cout << I_h_2h_matrix << std::endl;
        std::cout << ":::::::" << std::endl;
        weightedJacobi(A1, v1, f1, MG.v1);
        if (v1.size() <= 1)
        { // 最粗网格终止
            Eigen::SparseLU<Eigen::SparseMatrix<double>> solver;
            solver.compute(A1);
            v1 = solver.solve(f1);
            weightedJacobi(A1, v1, f1, MG.v2);
            return v1;
        }
        Eigen::VectorXd f2 = I_h_2h_matrix * (f1 - A1 * v1);
        Eigen::VectorXd v2 = Eigen::VectorXd::Zero(f2.size());
        Eigen::SparseMatrix<double> A2 = I_h_2h_matrix * A1 * I_2h_h_matrix;
        v2 = V_cycles(A2, v2, f2);
        v1 += I_2h_h_matrix * v2;
        weightedJacobi(A1, v1, f1, MG.v2);
        return v1;
    }

    Eigen::VectorXd FMG(Eigen::SparseMatrix<double> &A1, Eigen::VectorXd &f1)
    {
        Eigen::VectorXd v1 = Eigen::VectorXd::Zero(f1.size());
        const auto &I_h_2h_matrix = I_h_2h(f1.size() + 1);
        const auto &I_2h_h_matrix = I_2h_h(f1.size() + 1);
        // std::cout << "flag==" << flag << std::endl;
        // std::cout << "A1:" << std::endl;
        // std::cout << A1 << std::endl;
        // std::cout << "I_h_2h:" << std::endl;
        // std::cout << I_h_2h_matrix << std::endl;
        // std::cout << ":::::::" << std::endl;
        if (f1.size() <= 1)
        { // 最粗网格终止
            v1 = V_cycles(A1, v1, f1);
            return v1;
        }
        Eigen::VectorXd f2 = I_h_2h_matrix * (f1);
        Eigen::SparseMatrix<double> A2 = I_h_2h_matrix * A1 * I_2h_h_matrix;
        Eigen::VectorXd v2 = Eigen::VectorXd::Zero(f2.size());
        v2 = FMG(A2, f2);
        v1 = I_2h_h_matrix * v2;
        v1 = V_cycles(A1, v1, f1);
        return v1;
    }

    void TG()
    {
        std::cout << A << std::endl;
        std::cout << I_h_2h(m + 3) << std::endl;
        std::cout << I_2h_h(m + 3) << std::endl;
        std::cout << "n = " << n << std::endl;

        for (int i = 0; i < U.size(); ++i)
        {
            double x = i * h;
            std::cout << "x = " << x << ", u(x) = " << U[i] << std::endl;
        }

        for (int i = 0; i < MG.max_iterations; i++)
        {
            // Step 1: Apply weighted Jacobi smoothing
            weightedJacobi(A, U, f, MG.v1);

            // Step 2: Compute residual r_h = f - A * U
            Eigen::VectorXd r_h = f - A * U;

            // Step 3: Restrict residual to coarser grid
            const auto &I_h_2h_matrix = I_h_2h(m + 3);
            Eigen::VectorXd r_2h = I_h_2h_matrix * r_h;

            // Step 4: Compute coarse grid operator A_2h = I_h_2h * A * I_2h_h
            const auto &I_2h_h_matrix = I_2h_h(m + 3);
            Eigen::SparseMatrix<double> A_2h = I_h_2h_matrix * A * I_2h_h_matrix;

            // Step 5: Solve A_2h * e_2h = r_2h
            Eigen::SparseLU<Eigen::SparseMatrix<double>> solver_2h;
            solver_2h.compute(A_2h);
            if (solver_2h.info() != Eigen::Success)
            {
                std::cerr << "Error: Sparse matrix decomposition failed for A_2h.\n";
                exit(1);
            }

            Eigen::VectorXd e_2h = solver_2h.solve(r_2h);
            if (solver_2h.info() != Eigen::Success)
            {
                std::cerr << "Error: Sparse matrix solver failed for A_2h.\n";
                exit(1);
            }

            // Step 6: Interpolate error back to fine grid
            Eigen::VectorXd e_h = I_2h_h_matrix * e_2h;

            // Step 7: Update solution U += e_h
            U += e_h;

            // Step 8: Apply weighted Jacobi smoothing again
            weightedJacobi(A, U, f, MG.v2);

            double residual_norm = (f - A * U).norm() / f.norm();
            if (residual_norm < MG.relative_accuracy)
                break;
        }
        std::cout << "n = " << n << std::endl;
    }

    void solve()
    {
        if (MG.type == MGParams_One_Dimension::V_Cycles)
        {
            for (int i = 0; i < MG.max_iterations; i++)
            {
                U = V_cycles(A, U, f);
                double relative_residual = (f - A * U).norm() / f.norm();
                if (relative_residual <= MG.relative_accuracy)
                {
                    break; // 达到收敛条件，退出迭代
                }
            }
        }

        if (MG.type == MGParams_One_Dimension::FMG)
        {
            for (int i = 0; i < MG.max_iterations; i++)
            {
                U = FMG(A, f);
                double relative_residual = (f - A * U).norm() / f.norm();
                if (relative_residual <= MG.relative_accuracy)
                {
                    break; // 达到收敛条件，退出迭代
                }
            }
        }
        return;
    }
    void Solve()
    {
        Eigen::SparseLU<Eigen::SparseMatrix<double>> solver;
        solver.compute(A);
        if (solver.info() != Eigen::Success)
        {
            std::cerr << "Error: Sparse matrix decomposition failed.\n";
            exit(1);
        }

        U = solver.solve(f);
        if (solver.info() != Eigen::Success)
        {
            std::cerr << "Error: Sparse matrix solver failed.\n";
            exit(1);
        }
    }

    const Eigen::VectorXd &GetSolution() const
    {
        return U;
    }
};

#endif // TEST_HPP